import os
import cv2
import numpy as np
from functools import lru_cache
from moviepy.editor import VideoFileClip


def clip_video_to_frames(video_file, start_ms, end_ms,
                         ignore_unsuccess=False):
    """
    Clip video and return as a BGR frame list.

    Args:
        video_file (string): input video path.
        start_ms (float): starting time in milliseconds.
        end_ms (float): ending time in milliseconds. When its value is None,
            means that the duration.
    """
    clip = VideoFileClip(video_file)
    duration_ms = clip.duration * 1000
    if end_ms is None:
        end_ms = duration_ms

    end_ms = min(end_ms, duration_ms)
    new_clip = clip.subclip(start_ms / 1000., end_ms / 1000.)
    frames = [i[:, :, ::-1] for i in new_clip.iter_frames()]
    clip.close()
    new_clip.close()
    return frames


@lru_cache(maxsize=32)
def read_all_frames(path, return_id=False):
    idx, all_frames = 0, []
    img_file = os.path.join(path, '{}.jpg'.format(idx))
    while os.path.exists(img_file):
        if return_id:
            all_frames.append(str(idx))
        else:
            img = cv2.imread(img_file)
            all_frames.append(img)

        idx += 1
        img_file = os.path.join(path, '{}.jpg'.format(idx))

    return all_frames


def read_frames_dir_with_fids(path, frame_ids):
    frames = []
    for fid in frame_ids:
        img_file = os.path.join(path, '{}.jpg'.format(fid))
        img = cv2.imread(img_file)
        frames.append(img)
    return frames


def read_frames_dir(path, ts, te, length=5500, return_id=False):
    all_frames = read_all_frames(path, return_id=return_id)
    i = int(round(ts / (length + 0.0) * len(all_frames)))
    j = int(round(te / (length + 0.0) * len(all_frames)))
    return all_frames[i:j]


def get_video_length(video_file):
    """
    Get video length in milliseconds and frames count.

    Args:
        video_file (string): input video path.
    """
    clip = VideoFileClip(video_file)
    duration_ms = int(clip.duration * 1000)
    total_frames = int(clip.duration * clip.fps)
    clip.close()
    return duration_ms, total_frames


def get_video_extra_info(video_file):
    """
    Get video fps, width and height.
    This function is useful for getting config of VideoWriter.

    Args:
        video_file (string): input video path.
    """
    clip = VideoFileClip(video_file)
    fps, w, h = clip.fps, clip.w, clip.h
    clip.close()
    return fps, w, h


def convert_to_h264(video_file):
    """
    Convert video codec to H264.
    MP4V codec video generated by VideoWriter cannot be read by browser.
    Use this function to convert it via ffmpeg!
    """
    assert video_file.endswith('.mp4')
    tmp_video_file = ''.join(video_file.split('.')[:-1]) + '_tmp.mp4'
    try:
        os.system('ffmpeg -i {} -vcodec libx264 {}'.format(
            video_file, tmp_video_file))
        os.system('mv {} {}'.format(tmp_video_file, video_file))
        return True
    except Exception:
        return False


class VideoWriter(object):
    def __init__(self, filename, frame_size, fps):
        """
        Custom video writer.

        Args:
            filename (string): output video path.
            frame_size (tuple): (w, h) of the video frame.
            fps (float): FPS of the video.
        """
        self.frame_size = frame_size
        self.vout = cv2.VideoWriter()
        fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v')
        success = self.vout.open(filename, fourcc, fps, frame_size, True)
        if not success:
            raise RuntimeError('Error when create %s.' % filename)

    def add_frame(self, frame):
        assert frame.dtype == np.uint8
        assert frame.shape == (self.frame_size[1], self.frame_size[0], 3)
        self.vout.write(frame)

    def close(self):
        self.vout.release()
        del self.vout
